# -*- coding: utf-8 -*-
# File : inventory_alert.py
# Author: taoyahui
# Date : 2021/4/4
import sys
sys.path.append('..')

from utils.es_util import get_cve_alerts,get_distinct_ip
from conf import hostname, es_username, es_passwd
from utils import mysql_util


def inventory_alert(time_ago):
    conn = mysql_util.get_conn()
    cursor = mysql_util.get_cursor(conn)

    ips = get_distinct_ip(hostname, es_username, es_passwd)

    for ip in ips:
        cve_detail_res_arr = get_cve_alerts(hostname, es_username, es_passwd, ip, time_ago)
        inventory_cve_dicts = []
        for res in cve_detail_res_arr:
            try:
                source = res['hits']['hits'][0]['_source']
                agent = source['agent']
                name = agent['name']
                ip = agent['ip']
                vul = source['data']['vulnerability']
                cve = vul['cve']
                severity = vul['severity']
                package_condition = vul['package']['condition']
                package_name = vul['package']['name']
                package_version = vul['package']['version']
                package_arch = vul['package']['architecture']
                score = float(res['hits']['hits'][0]['_score'])
                title = vul['title']
                advisories_ids = None
                if 'advisories_ids' in vul.keys():
                    advisories_ids = ','.join(vul['advisories_ids'])
                sql = f"select count(1) from cve where cve_id = '{cve}'"
                if mysql_util.exec_sql(conn, cursor, sql)[0][0] == 0:
                    insert_sql = f"insert into cve(cve_id, severity, score, pac_condition, pac_name, pac_version, pac_arch, title, advisories_ids) " \
                                 f"values ('{cve}', '{severity}', {score}, '{package_condition}'," \
                                 f"'{package_name}','{package_version}','{package_arch}', '{title}', '{advisories_ids}')"
                    # print(insert_sql)
                    mysql_util.exec_sql(conn, cursor, insert_sql)
                mdict = {'ip': ip, 'cve_id': cve}
                inventory_cve_dicts.append(mdict)
            except Exception as e:
                print(e)
        remove_sql = f"delete from inventory_cve where inventory_ip = '{ip}'"
        mysql_util.exec_sql(conn, cursor, remove_sql)
        for mdict in inventory_cve_dicts:
            relate_insert_sql = f"insert into inventory_cve(inventory_ip, cve_id) values ('{mdict['ip']}', '{mdict['cve_id']}')"
            mysql_util.exec_sql(conn, cursor, relate_insert_sql)


        # print(alerts_resp)

    # 查找出所有资产

if __name__ == '__main__':
    if sys.argv.__len__() > 1:
        time_ago = sys.argv[1]
    else:
        time_ago = '12h'
    inventory_alert(time_ago)
